package graph

import (
	"container/heap"
	"math"
)

type Pair struct {
	Name     string
	Distance int
}

type PQ []Pair

func (p PQ) Len() int {
	return len(p)
}

func (p PQ) Less(i, j int) bool {
	return p[i].Distance < p[j].Distance
}

func (p PQ) Swap(i, j int) {
	p[i], p[j] = p[j], p[i]
}

func (p *PQ) Push(x interface{}) {
	*p = append(*p, x.(Pair))
}

func (p *PQ) Pop() (v interface{}) {
	*p, v = (*p)[:p.Len()-1], (*p)[p.Len()-1]
	return
}

func Dijkstra(graph map[string][]Pair, s string) (map[string]string, map[string]int) {
	pq := new(PQ)

	// 这两个数组（map) + 优先对列 是dijkstra算法的核心
	prevs := make(map[string]string)
	dists := make(map[string]int)

	visited := make(map[string]bool)
	heap.Init(pq)
	heap.Push(pq, Pair{
		Name:     s,
		Distance: 0,
	})

	// init distance
	for k := range graph {
		if k == s {
			dists[k] = 0
			prevs[k] = "x"
		} else {
			dists[k] = math.MaxInt64
		}
	}

	for {
		if pq.Len() <= 0 {
			break
		}
		pair := heap.Pop(pq).(Pair)
		vertex := pair.Name
		dist := pair.Distance

		visited[vertex] = true
		for _, v := range graph[vertex] {
			if !visited[v.Name] {
				if dist+v.Distance < dists[v.Name] {
					heap.Push(pq, Pair{
						Name:     v.Name,
						Distance: dist + v.Distance,
					})
					prevs[v.Name] = vertex
					dists[v.Name] = dist + v.Distance
				}
			}
		}
	}

	return prevs, dists

}
